
#include <json/json.h>

#include "base_util/utils.h"
// #include "ocr/ocr_det.h"
// #include "ocr/ocr_rec.h"
// #include "cls/cls_predictor.h"
#include "det/det_predictor.h"
// #include "seg/seg_predictor.h"


#include "common/model_config.h"
// #include "ocr/postprocess_op.h"

#include "paddle/pdmodel_manager.h"


namespace ai {

PDModelManager::PDModelManager(LogInfo *log_info):
  ai::ModelManager(log_info), init_ok(-1),
  stop_flag1(common::CameraClosed), stop_flag2(common::CameraClosed) {

  // cls_pd_ = new ClsPredictor(log_ifo);
  // text_det_ = new OCRDetector(log_ifo);
  // text_rec_ = new OCRRecognizer(log_ifo);
  det_pd_ = new PDDetPredictor(log_ifo);
  // seg_pd_ = new SegPredictor(log_ifo);
  // ocr_pp = new OCRPostProcessor();

}

PDModelManager::~PDModelManager(){

}

int PDModelManager::init_model(std::vector<InitModelData> &init_model_datas)
{
  int ret = 0;
  for (int idx = 0; idx < init_model_datas.size(); idx++) {
    ModelConfig *model_config = new ModelConfig(log_ifo);
    paddle_infer::Config config;
    auto& imd = init_model_datas[idx];
    ret = load_model(imd.model_dir, 
                    imd.model_type,
                    model_config, config);
    if (ret !=0){return ret;}

    predictors[imd.model_ind] = std::move(paddle::CreatePaddlePredictor(config));
    configs[imd.model_ind] = std::move(model_config);
    // 在初始化时存储模型信息
    total_mdl[imd.model_name] = ModuleInfo(imd.model_ind, imd.model_type);
  }

  // log
  std::string log_str = "total model:\n";
  for (auto& tm : total_mdl) {
    log_str += "\t\t" + tm.first + "\t\t" + std::to_string(tm.second.mdl_idx)+ "\n";
  }
  spdlog::get("logger")->info("PDModelManager::init_model: {} ", log_str);

  return 0;
}

// 未实现
int PDModelManager::stop_run_thread() {}
int PDModelManager::start_run_thread(int cam_size) {}
int PDModelManager::set_data(int idx, cv::Mat& img) {}
int PDModelManager::set_data(int idx, stream::ImageBlob& img) {return -1;}
int PDModelManager::get_data(int idx, std::vector<BaseInfo*>& data) {}

int PDModelManager::rec_img(std::string algo_type,
                      cv::Mat& img, 
                      std::vector<BaseInfo*> &data_info) {
  // rec_mut.lock();
	int ret = 0;
  if (total_mdl[algo_type].algoType == model_cls) {
    ret = this->cls(img, data_info, total_mdl[algo_type].mdl_idx);
  } else if (total_mdl[algo_type].algoType == model_det) {
    ret = this->det(img, data_info, total_mdl[algo_type].mdl_idx);
  } else if (total_mdl[algo_type].algoType == model_text_rec) {
    ret = this->text_rec_lite(img, data_info, total_mdl[algo_type].mdl_idx);
  } else if (total_mdl[algo_type].algoType == model_seg) {
    ret = this->seg(img, data_info, total_mdl[algo_type].mdl_idx);
  } else if (total_mdl[algo_type].algoType == model_plate_rec) {
    ret = this->plate_rec(img, data_info, total_mdl[algo_type].mdl_idx);
  } else if (total_mdl[algo_type].algoType == model_none) {
    ret = 0;
  } else {
    spdlog::get("logger")->info("rec_img eror...");
  }
  // rec_mut.unlock();
	return ret;
}

int PDModelManager::plate_rec(cv::Mat& img, 
                  std::vector<BaseInfo*> &rec_ifo, 
                  int mdl_idx ) {
  int ret = 0;
  // std::cout << "plate_rec..." << ret << std::endl;
  // det
  std::vector < BaseInfo* > text_lines;
  ret = this->text_detect(img, text_lines, mdl_idx);
  if (ret != 0) { return ret;}
  // std::cout << "text_detect ret: " << ret << std::endl;
  // std::cout << "text_lines.size(): " << text_lines.size() << std::endl;
  // 车牌识别 只保留一个

  int max_v = 0;
  int max_idx = -1;
  for (int idx = 0; idx < text_lines.size(); idx++) {
    TextDetInfo* difo = (TextDetInfo*)(text_lines[idx]);
    int area = cv::contourArea(difo->four_points);
    if (area > max_v) { max_v = area; max_idx = idx; }
  }
  if (max_idx < 0) { return -1;}

  // rec
  // std::vector<TextRecInfo > ocr_infos;
  TextDetInfo* difo = (TextDetInfo*)(text_lines[max_idx]);
  TextRecInfo crnn_info;
  ret = this->text_recognise(img, *difo, crnn_info, mdl_idx+1);
  if (ret != 0){
    std::cout << "text_recognise: " << ret << std::endl;
    return ret;;
  }
  rec_ifo.push_back(new TextRecInfo(difo->score, difo->rrect, difo->four_points, crnn_info));
  // std::cout << crnn_info.text << std::endl;
  // ocr_infos.push_back(crnn_info);

  for (auto& tl : text_lines) { delete tl;}
  return ret;
}


int PDModelManager::text_recognise(cv::Mat& img, 
                  std::vector<BaseInfo*> &rec_ifo, 
                  int mdl_idx ) {
  int ret = 0;
  // det
  std::vector < BaseInfo* > text_lines;
  ret = this->text_detect(img, text_lines, mdl_idx);
  if (ret != 0) { return ret;}
  // std::cout << "text_detect ret: " << ret << std::endl;
  // std::cout << "text_lines.size(): " << text_lines.size() << std::endl;

  // rec
  for (int idx = text_lines.size()-1; idx >= 0 ; idx--) {
    TextDetInfo* difo = (TextDetInfo*)(text_lines[idx]);
    TextRecInfo crnn_info;
    ret = this->text_recognise(img, *difo, crnn_info, mdl_idx+1);
    if (ret != 0){
      std::cout << "text_recognise: " << ret << std::endl;
      continue;
    }
    rec_ifo.push_back(new TextRecInfo(difo->score, difo->rrect, difo->four_points, crnn_info));
    // std::cout << crnn_info.text << std::endl;
    if(log_ifo->log_level_2){spdlog::get("logger")->info("{}", crnn_info.text); }
  }
  if(log_ifo->log_level_2){spdlog::get("logger")->info("\n\n\n"); }

  for (auto& tl : text_lines) { delete tl;}
  return ret;
}


int PDModelManager::free_model()
{
  for (auto it = predictors.begin(); it != predictors.end(); it++) {
    it->second.release();
  }
  return 0;
}

int PDModelManager::load_model(const std::string &model_dir, 
                              ModelType mt,
                              ModelConfig* model_config,
                              paddle_infer::Config &config) 
{
  int ret = 0;
  if ((ret = model_config->load_config(model_dir)) != 0){ return ret; }
  if (model_config->algorithm != mt) { return model_model_type_not_same_error;}

  config.SetModel(model_dir + "/model", model_dir + "/params");
  if (model_config->use_gpu) {
    config.EnableUseGpu(model_config->gpu_mem, model_config->gpu_id);
    std::map<std::string, paddle::AnalysisConfig::Precision> TRT_MAP = {
        {"fp32", paddle::AnalysisConfig::Precision::kFloat32},
        {"fp16", paddle::AnalysisConfig::Precision::kHalf},
        {"int8", paddle::AnalysisConfig::Precision::kInt8}
    };
    if (TRT_MAP.find(model_config->predictor_type) != TRT_MAP.end()) {
        auto precision = TRT_MAP[model_config->predictor_type];
        bool use_cab = (precision == paddle::AnalysisConfig::Precision::kInt8);
        config.EnableTensorRtEngine(1 << 30, model_config->batch_size, model_config->min_subgraph_size,
            precision, true, false);
        // int imgH = this->rec_image_shape_[1];
        // int imgW = this->rec_image_shape_[2];
        // std::map<std::string, std::vector<int>> min_input_shape = {
        //     {"x", {1, 3, imgH, 10}}, {"lstm_0.tmp_0", {10, 1, 96}}};
        // std::map<std::string, std::vector<int>> max_input_shape = {
        //     {"x", {this->rec_batch_num_, 3, imgH, 2500}},
        //     {"lstm_0.tmp_0", {1000, 1, 96}}};
        // std::map<std::string, std::vector<int>> opt_input_shape = {
        //     {"x", {this->rec_batch_num_, 3, imgH, imgW}},
        //     {"lstm_0.tmp_0", {25, 1, 96}}};

        // config.SetTRTDynamicShapeInfo(min_input_shape, max_input_shape,
        //                               opt_input_shape);
    }
  } else {
    config.DisableGpu();
    if (model_config->use_mkldnn) { 
      config.EnableMKLDNN(); 
      config.SetMkldnnCacheCapacity(10);
    }
    config.SetCpuMathLibraryNumThreads(model_config->cpu_threads);
  }

  if (model_config->algorithm == model_text_rec) { 
    auto pass_builder = config.pass_builder();
    pass_builder->DeletePass("matmul_transpose_reshape_fuse_pass");
  }

  // false for zero copy tensor
  config.SwitchUseFeedFetchOps(false);
  // true for multiple input
  config.SwitchSpecifyInputNames(true);
  config.SwitchIrOptim(true);
  config.EnableMemoryOptim();
  config.DisableGlogInfo();
  return 0;
}


int PDModelManager::text_detect(cv::Mat& img, 
                                std::vector <BaseInfo*> &text_lines, 
                                int mdl_idx,
                                bool sort_by_default) {
  // int ret =0;
  // ret = text_det_->Run(img, text_lines, this, mdl_idx);
  // if (ret !=0 ){return ret;}
  // if (text_lines.size() == 0){return model_not_detect_text_lines;}
  // if (log_ifo->log_all_img){
  //   cv::Mat save_image;
  //   visualize_db(img,save_image,text_lines, false);
  //   std::string save_path = bf::path(log_ifo->get_log_path()).append(std::to_string(log_ifo->counter()) + "_db_det.jpg").string();
  //   cv::imwrite(save_path, save_image);
  // }

  // // for(int n = 0; n < text_lines.size(); n++){
  // //   TextDetInfo* tl = (TextDetInfo*)text_lines[n];
  // //   // four_points 已经排序, 顺时针方向
  // //   tl->proposal_points.insert(tl->proposal_points.begin(), tl->four_points);
  // //   // 统一放到 rrects 里
  // //   for (auto ps : tl->proposal_points){
  // //     cv::RotatedRect rt_temp = cv::minAreaRect(ps);
  // //     tl->proposal_rrects.push_back(rt_temp);
  // //   }
  // //   tl->rrect = tl->proposal_rrects[0];
  // // }

  // // // sort box
  // // // float mean_height = 0.;
  // // // for(auto bi : text_lines){mean_height += bi.rrect.size.height;}
  // // // mean_height = mean_height * 1.0 / text_lines.size();
  // // if (!sort_by_default) {return 0;}
	// // std::sort(text_lines.begin(), text_lines.end(), 
  // //   [&](BaseInfo* aa, BaseInfo* bb){
  // //     TextDetInfo* a = (TextDetInfo*)aa;
  // //     TextDetInfo* b = (TextDetInfo*)bb;
  // //     int ay = (a->four_points[0].y + a->four_points[1].y + a->four_points[2].y + a->four_points[3].y) / 4;
  // //     int ax = (a->four_points[0].x + a->four_points[1].x + a->four_points[2].x + a->four_points[3].x) / 4;
  // //     int by = (b->four_points[0].y + b->four_points[1].y + b->four_points[2].y + b->four_points[3].y) / 4;
  // //     int bx = (b->four_points[0].x + b->four_points[1].x + b->four_points[2].x + b->four_points[3].x) / 4;
  // //     if (std::abs(ay - by) < 5){ return ax < bx;} 
  // //     else { return ay < by; }
  // // });
  return 0;
}

int PDModelManager::text_rec_lite(cv::Mat& img, 
                                  std::vector<BaseInfo*> &rec_ifo, 
                                  int mdl_idx )
{
  int ret = 0;
  // ret = text_rec_->Run(img, rec_ifo, this, mdl_idx);
  return ret;
}

int PDModelManager::text_recognise(cv::Mat& img, 
                                  TextDetInfo &det_info,
                                  TextRecInfo &crnn_info, 
                                  int model_idx,
                                  float unclip_ratio)
{
  int ret = 0;
  cv::Mat crop_image;

  // GetRotateCropImage(img, det_info.four_points, crop_image);
  // if (log_ifo->log_all_img){
  //   std::string save_path = bf::path(log_ifo->get_log_path()).append(std::to_string(log_ifo->counter()) + "__crop_image.jpg").string();
  //   cv::imwrite(save_path, crop_image);
  // }
  // std::vector<BaseInfo*> rec_ifo;
  // ret = text_rec_->Run(crop_image, rec_ifo, this, model_idx);
  // if (ret != 0) { 
  //   std::cout << "text_rec_->Run: " << ret << std::endl;
  //   return ret;
  // }
  // TextRecInfo c_info = *(TextRecInfo*)(rec_ifo[0]);
  // c_info.four_points = det_info.four_points;
  // for (auto& ri : rec_ifo) { delete ri;}
  // crnn_info = c_info;

  // // std::vector<TextRecInfo> crnn_infos;
  // // for (int i=0;i<det_info.proposal_points.size();i++){
  // //   cv::Mat crop_image;

  // //   std::vector<cv::Point> new_points = det_info.proposal_points[i];
  // //   if (unclip_ratio > 0.1){
  // //     std::vector<cv::Point> four_points;
  // //     float score;
  // //     if ((ret = ocr_pp->unclip_for_points(new_points, unclip_ratio,
  // //                                   0.1, *img, *img, four_points,
  // //                                   score)) !=0 ){
  // //       std::cout << "unclip_for_points: " << ret << std::endl;
  // //       continue;
  // //     }
  // //     new_points = four_points;
  // //   }
  // //   for (auto& p : new_points) {
  // //     if (p.x < 0 || p.y < 0 || p.x > img.cols || p.y > img.rows) { 
  // //       std::cout << "new_points no ok..." << std::endl;
  // //       continue;
  // //     }
  // //   }
  // //   GetRotateCropImage(*img, new_points, crop_image);
  // //   if (log_ifo->log_all_img){
  // //     std::string save_path = bf::path(log_ifo->get_log_path()).append(std::to_string(log_ifo->counter()) + "__crop_image.jpg").string();
  // //     cv::imwrite(save_path, crop_image);
  // //   }
  // //   std::vector<BaseInfo*> rec_ifo;
  // //   ret = text_rec_->Run(&crop_image, rec_ifo, this, model_idx);
  // //   if (ret != 0) { 
  // //     std::cout << "text_rec_->Run: " << ret << std::endl;
  // //     continue;;
  // //   }
  // //   TextRecInfo c_info = *(TextRecInfo*)(rec_ifo[0]);
  // //   for (auto& ri : rec_ifo) { delete ri;}


  // //   crnn_infos.push_back(TextRecInfo(det_info.score, 
  // //                         det_info.proposal_rrects[i], 
  // //                         det_info.proposal_points[i], c_info));
  // // }

  // // int most_suitable_idx=0;
  // // if (crnn_infos.size() == 0){
  // //   std::cout << "crnn_infos.size() == 0" << std::endl;
  // //   return model_nothing_to_recognise;
  // // }
  // // else if (crnn_infos.size() > 1){
  // //   std::vector<std::string> total_strings;
  // //   std::vector<int> string_index;
  // //   std::vector<int> crnn_data_info_index;
  // //   for (int i =0;i<crnn_infos.size();i++){
  // //     std::vector<std::string>::iterator it= total_strings.end();
  // //     it = std::find_if(total_strings.begin(), total_strings.end(), [&](std::string &a){
  // //       return crnn_infos[i].text == a;
  // //     });
  // //     if (it == total_strings.end()){
  // //       total_strings.push_back(crnn_infos[i].text);
  // //       string_index.push_back(total_strings.size() -1);
  // //     }else{
  // //       string_index.push_back(it - total_strings.begin());
  // //     }
  // //     crnn_data_info_index.push_back(i);
  // //   }

  // //   int max_string_index = find_mode(string_index);
  // //   for (int ii=0;ii<string_index.size();ii++){
  // //     if (string_index[ii] == max_string_index){
  // //       most_suitable_idx = crnn_data_info_index[ii];
  // //       break;
  // //     }
  // //   }
  // // }

  // // crnn_info = crnn_infos[most_suitable_idx];
  return ret;
}

int PDModelManager::cls(cv::Mat& img,
                std::vector<BaseInfo*>& cls_info, 
                int md_idx) { 
  // return cls_pd_->Run(img, cls_info, this, md_idx);
}
int PDModelManager::cls(std::vector<cv::Mat> &imgs,
                std::vector<BaseInfo*>& cls_ifos, 
                int md_idx) { 
  // return cls_pd_->Run(imgs, cls_ifos, this, md_idx);
}
int PDModelManager::seg(cv::Mat& img, 
            std::vector<BaseInfo*> &seg_info,
            int mdl_idx)
{
  // return seg_pd_->Run(img, seg_info, this, mdl_idx);
}

int PDModelManager::det(cv::Mat& img, 
            std::vector<BaseInfo*> &det_info,
            int model_idx)
{
  return det_pd_->Run(img, det_info, this, model_idx);
}
int PDModelManager::det(std::vector<cv::Mat> &imgs,
              std::vector<std::vector<DetInfo>>& det_infos, 
              int model_idx)
{
  // return det_pd_->Run(imgs, det_infos, this, model_idx);
  return 0;
}

} // namespace paddle_infer